# To change this template, choose Tools | Templates
# and open the template in the editor.
import telnetlib
import getopt
import sys

__author__="winstonlee"
__date__ ="$Oct 3, 2010 7:56:47 PM$"

class BillionModem:

    def __init__(self):
        self.verbose = False

"""
The base Billion Telnet class.

Currently assumes that:
* models are quite consistent in Telnet commands and responses
* firmware/software version doesn't really matter
"""
class BillionModemTelnet(BillionModem):
    # The telnet prompt expectations
    PROMPT_LOGIN = "Login: "
    PROMPT_PASSWORD = "Password: "
    PROMPT_LOGGED_IN = ">"

    # The telnet commands
    COMMAND_PORT_A1_STATUS = "port a1 status\r\n"
    COMMAND_LOGOUT = "user logout\r\n"
    COMMAND_PPPOE_SET_TRANSPORT_WANLINK_ENABLED = "pppoe set transport wanlink enabled\r\n"
    COMMAND_PPPOE_SET_TRANSPORT_WANLINK_DISABLED = "pppoe set transport wanlink disabled\r\n"
    COMMAND_SYSTEM_RESTART = "system restart\r\n"
    COMMAND_SYSTEM_INFO = "system info\r\n"

    def __init__(self, host, username, password):
        BillionModem.__init__(self)
        self.host = host
        self.username = username
        self.password = password
        self.telnet = telnetlib.Telnet()
        #Override the prompt string
        BillionModemTelnet.PROMPT_LOGGED_IN = "%s>" % (username,)

    def connect(self):
        if self.verbose: print "Connecting as %s to %s" % (self.username, self.host)
        self.telnet.open(self.host, 23)
        self.telnet.read_until(BillionModemTelnet.PROMPT_LOGIN, 5)
        self.telnet.write(self.username + "\r\n")
        self.telnet.read_until(BillionModemTelnet.PROMPT_PASSWORD, 5)
        self.telnet.write(self.password + "\r\n")
        #Login failed
        logged_in = self.telnet.expect(['Login failed', 'Login successful'], 5)
        if self.verbose: print "Login result %s" % (logged_in,)
        if (0 == logged_in[0]):
            # Authentication failed - logged_in[0] corresponds to the match in the expectation's list
            if self.verbose: print "Bad username or password"
            self.telnet.close()
            raise ValueError("Bad username or password")
        elif (1 == logged_in[0]):
            if self.verbose: print "Connected"
            self.telnet.read_until(BillionModemTelnet.PROMPT_LOGGED_IN)
        else:
            if self.verbose: print "Unknown state - how did you do this? %s" % (logged_in,)
            raise ValueError("Unknown state - check username/password")

    def restart(self):
        if self.verbose: print "Sending command %s" % (BillionModemTelnet.COMMAND_SYSTEM_RESTART,)
        self.telnet.write(BillionModemTelnet.COMMAND_SYSTEM_RESTART)

    def logout(self):
        if self.verbose: print "Logging user out"
        self.telnet.write(BillionModemTelnet.COMMAND_LOGOUT)
        self.telnet.close()
        if self.verbose: print "Logged out"

    def is_wan_connected(self):
        if self.verbose: print "Sending command %s" % (BillionModemTelnet.COMMAND_PORT_A1_STATUS,)
        self.telnet.write(BillionModemTelnet.COMMAND_PORT_A1_STATUS)
        result_string = self.telnet.read_until(BillionModemTelnet.PROMPT_LOGGED_IN, 15)
        if self.verbose: print "Response is %s" % (result_string,)
        result_map = self.__parse_wan_conected(result_string)
        if self.verbose: print "Parsed response is %s" % (result_map,)
        if self.verbose: print "Connection status is %s" % (result_map["Connected"],)
        return (result_map["Connected"])

    def get_system_info(self):
        if self.verbose: print "Sending command %s" % (BillionModemTelnet.COMMAND_SYSTEM_INFO,)
        self.telnet.write(BillionModemTelnet.COMMAND_SYSTEM_INFO)
        result_string = self.telnet.read_until(BillionModemTelnet.PROMPT_LOGGED_IN, 15)
        if self.verbose: print "Response is %s" % (result_string,)
        result_map = self.__parse_system_info(result_string)
        if self.verbose: print "Response is %s" % (result_map,)
        return result_map

    def __parse_system_info(self, result):
        new_res = self.__strip_logged_in_prompt(result)
        new_res = self.__strip_last_command(new_res, BillionModemTelnet.COMMAND_SYSTEM_INFO)
        new_res = new_res.replace("Global System Configuration:", "")
        result_list = new_res.split("\r\n")
        result_map = {}
        for item in result_list:
            if '' != item:
                x = item.split(": ")
                result_map[x[0].strip()] = x[1].strip()
        return result_map

    def __parse_wan_conected(self, result):
        new_res = self.__strip_logged_in_prompt(result)
        new_res = self.__strip_last_command(new_res, BillionModemTelnet.COMMAND_PORT_A1_STATUS)
        result_list = new_res.split("\r\n")
        result_map = {}
        for item in result_list:
            x = item.split(" = ")
            result_map[x[0].strip()] = x[1]
        return result_map

    def __strip_last_command(self, result, command):
        return result.lstrip(command)

    def __strip_logged_in_prompt(self, result):
        return result.rstrip('\r\n' + BillionModemTelnet.PROMPT_LOGGED_IN)

    def connect_wan(self):
        if self.verbose: print "Sending command %s" % (BillionModemTelnet.COMMAND_PPPOE_SET_TRANSPORT_WANLINK_DISABLED,)
        self.telnet.write(BillionModemTelnet.COMMAND_PPPOE_SET_TRANSPORT_WANLINK_DISABLED)
        self.telnet.read_until(BillionModemTelnet.PROMPT_LOGGED_IN, 15)
        if self.verbose: print "Sending command %s" % (BillionModemTelnet.COMMAND_PPPOE_SET_TRANSPORT_WANLINK_ENABLED,)
        self.telnet.write(BillionModemTelnet.COMMAND_PPPOE_SET_TRANSPORT_WANLINK_ENABLED)
        self.telnet.read_until(BillionModemTelnet.PROMPT_LOGGED_IN, 15)

"""
The Billion 7404VGP-M BillionModemTelnet implementation
"""
class Billion7404VGPM(BillionModemTelnet):

    def __init__(self, host, username, password):
        BillionModemTelnet.__init__(self, host, username, password)

##################
## The commands ##
##################

"""
Forces a reconnection. This is the same as disconnecting and reconnecting.
"""
def reconnect(telnet):
    telnet.connect()
    telnet.connect_wan()
    telnet.logout()

"""
Checks the wanlink status and forces a reconnection if the status is not 'true'
"""
def reconnect_if_down(telnet):
    telnet.connect()
    if ("true" != telnet.is_wan_connected()):
        if telnet.verbose: print "Connection is down, forcing reconnection"
        telnet.connect_wan()
    else:
        if telnet.verbose: print "Connection is up, nothing to do"
    telnet.logout()

"""
Queries the modem for its system information and parses the result.
"""
def sys_info(telnet):
    telnet.connect()
    dict = telnet.get_system_info()
    for a, v in dict.iteritems():
        print "%s = %s" % (a, v)
    telnet.logout()

def restart(telnet):
    telnet.connect()
    telnet.restart()
    # The logout could potentially fail but we don't really care
    try:
        telnet.logout()
    except:
        pass # swallow the exception since we are restarting the modem

"""
Instantiates the appropriate Telnet client/parser for the given model.
Uses the default model if none are matched.
"""
def _instantiate_telnet(host, username, password, model=None, verbose=False):
    if verbose: print "Attempting to instantiate %s model" % (model,)
    telnet = None
    try:
        telnet = globals()[model](host, username, password)
    except:
        if None == telnet:
            if verbose: print "Could not instantiate a client for the given model - using default"
            telnet = BillionModemTelnet(host, username, password)
            telnet.verbose = verbose

    return telnet

def usage():
    print """DESCRIPTION
    A Billion router helper

SYNOPSIS
    billion.py [-hv] [-i host_ip] [-u username] [-p password] [-m model] <COMMAND>

SWITCHES
    -i host_ip
        The host IP (generally 192.168.1.254)
        
    -u username
        The username of a user with the admin role

    -p password
        The password for the user

    -m model
        The model of Billion Modem - Billion7404VGPM

    -h --help
        This uage help

    -v --verbose
        Verbose output

COMMANDS
    restart
        restarts the modem (same as pressing reset)

    reconnect
        forces a reconnect on the wanlink

    reconnect-if-down
        checks the wanlink status and reconnects if necessary

    sys-info
        outputs the system information

EXAMPLE
    billion.py -i 192.168.1.254 -u admin -p admin reconnect
        forces a reconnect on the wanlink"""

def main():
    try:
        opts, args = getopt.getopt(sys.argv[1:], "i:u:p:m:vh", ["help", "verbose"])
    except getopt.GetoptError, err:
        print str(err)
        usage()
        sys.exit(2)

    verbose = False
    host = None
    username = None
    password = None
    model = None
    
    for o, a in opts:
        if o in ("--help", "-h"):
            usage()
            sys.exit()
        elif o in ("-i"):
            host = a
        elif o in ("-u"):
            username = a
        elif o in ("-p"):
            password = a
        elif o in ("-m"):
            model = a
        elif o in ("--verbose", "-v"):
            verbose = True
            print "options = " + str(opts)
            print "arguments = " + str(args)

    if host == None or username == None or password == None:
        print "Host, username and password are required"
        usage()
        sys.exit(2)

    # Check that there is only one command
    if len(args) > 1:
        print "Too many commands %s" % (str(args), )
        sys.exit(2)

    # Instantiate the appropriate client
    telnet = _instantiate_telnet(host, username, password, model, verbose)

    # Execute the given command
    command = args[0]
    if "reconnect" == command:
        if verbose: print "Execute reconnect command"
        reconnect(telnet)
    elif "reconnect-if-down" == command:
        if verbose: print "Execute reconnect-if-down command"
        reconnect_if_down(telnet)
    elif "sys-info" == command:
        if verbose: print "Execute sys-info command"
        sys_info(telnet)
    elif "restart" == command:
        if verbose: print "Execute restart command"
        restart(telnet)
    else:
        print "Unknown command %s" % (command, )
        exit(2)

if __name__ == "__main__":
    try:
        main()
    except Exception, e:
        print "ERR: %s " % (str(e),)
        exit(1)